package projectUtils;

import java.awt.Color;
import java.awt.image.BufferedImage;
import java.util.Random;

import featureExtraction.TokenCollection;
import featureExtraction.TokenDetection;
import featureExtraction.TokenRaw;

public class ImageDistortionUtil {
	
	private static Random randomGen = new Random();
	
	public static BufferedImage DistortImage(BufferedImage image, Double shiftAmount, Double noiseAmount, Double stretchPercentageStDev)
	{
		//return RandomErode(RandomDilate(AddNoise(RandomShift(image, shiftAmount), noiseAmount), scaleAmount), scaleAmount);
//		return AddNoise(RandomShift(image, shiftAmount), noiseAmount);
		return AddNoise(RandomShift(Stretch(image, stretchPercentageStDev), shiftAmount), noiseAmount);
		//return AddNoise(image, noiseAmount);
	}
	
	public static BufferedImage DistortImage(BufferedImage image)
	{
		image = Stretch(image, 0.1);
		image = RandomRotate(image, 0.05);
		image = Stretch(image, 0.1);
//		image = RandomDilate(image, 0.02);
//		image = RandomErode(image, 0.02);
//		image = RandomShift(image, 0.1);
//		image = AddNoise(image, 10.0);

		return image;
	}
	
	public static BufferedImage Undistort(BufferedImage image)
	{
		TokenDetection tokenDetector = new TokenDetection();
		TokenCollection tokenCollection = tokenDetector.SegmentImageTokens(image);
		double maxSize = 0;
		TokenRaw digitToken = null;
		for (TokenRaw token : tokenCollection.Tokens)
		{
			if ((token.X1 - token.X0) * (token.Y1 - token.Y0) > maxSize)
			{
				digitToken = token;
				maxSize = (token.X1 - token.X0) * (token.Y1 - token.Y0);
			}
		}

		image = ImageArrayUtil.CreateSubImage(image, digitToken.X0, digitToken.X1, digitToken.Y0, digitToken.Y1);
//		byte[][] se = new byte[3][3];
//		for (int k = 0; k < 3; k++)
//		{
//			for (int l = 0; l < 3; l++)
//				se[k][l] = 1;
//		}
//		for (int k = 0; k < 3; k++)
//		{
//			image = ImageArrayUtil.PadImage(image, 2, 2);
//			image =  ImageArrayUtil.DilateImage(image, se);
//		}
		image = ImageArrayUtil.PadImage(image, (int)(image.getWidth() * 0.4), (int)(image.getHeight() * 0.25));
		
		image = ImageArrayUtil.ArrayToImage(ImageArrayUtil.ShiftBasedOnCenterOfMass(ImageArrayUtil.ImageToArray(image)));
		return image;
	}
	
	public static BufferedImage RandomShift(BufferedImage image, Double percentageAmount)
	{
		Integer xShift = (int) (((double)image.getWidth() * percentageAmount) * (randomGen.nextDouble()));
		Integer yShift = (int) (((double)image.getHeight() * percentageAmount) * (randomGen.nextDouble()));
		BufferedImage returnImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_INT_RGB);
		for (int i = 0; i < image.getWidth(); i++)
		{
			for (int j = 0; j < image.getHeight(); j++)
			{
				Integer newX = i + xShift;
				Integer newY = j + yShift;
				if (newX < 0 || newX >= image.getWidth() || newY < 0 || newY >= image.getHeight())
					returnImage.setRGB(i, j, 0xFFFFFFFF);
				else
					returnImage.setRGB(i, j, image.getRGB(newX, newY));
			}
		}
		return returnImage;
	}
	
	public static BufferedImage AddNoise(BufferedImage image, Double intensityStDev)
	{
		BufferedImage returnImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_INT_RGB);
		for (int i = 0; i < image.getWidth(); i++)
		{
			for (int j = 0; j < image.getHeight(); j++)
			{
				if (image.getRGB(i, j) != 0xFFFFFFFF)
				{
					int intensity = ColorUtil.RGBToGreyScale(image.getRGB(i, j));
					intensity = Math.min(Math.max(0, intensity + (int)(randomGen.nextGaussian() * intensityStDev)), 255);
					returnImage.setRGB(i, j, new Color(intensity, intensity, intensity).getRGB());
				}
				else
					returnImage.setRGB(i, j, 0xFFFFFFFF);
			}
		}
		return returnImage;
	}
	
	public static BufferedImage Stretch(BufferedImage image, Double stretchPercentage)
	{
		double widthMultiplier = 1 + randomGen.nextDouble() * stretchPercentage;
		double heightMultiplier = 1 + randomGen.nextDouble() * stretchPercentage;
		if (randomGen.nextBoolean() == true)
			widthMultiplier = 1.0 / widthMultiplier;
		if (randomGen.nextBoolean() == true)
			heightMultiplier = 1.0 / heightMultiplier;
		int newWidth = (int) ((widthMultiplier) * image.getWidth());
		int newHeight = (int) ((heightMultiplier) * image.getHeight());
		BufferedImage scaledImage = ImageArrayUtil.CreateResizedCopy(image, newWidth, newHeight);
		BufferedImage returnImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_INT_RGB);
		for (int i = 0; i < image.getWidth(); i++)
		{
			for (int j = 0; j < image.getHeight(); j++)
			{
				int x = i + (newWidth - image.getWidth()) / 2;
				int y = j + (newHeight - image.getHeight()) / 2;
				if (x < 0 || y < 0 || x >= newWidth || y >= newHeight)
					returnImage.setRGB(i, j, 0xFFFFFFFF);
				else
					returnImage.setRGB(i, j, scaledImage.getRGB(x, y));
			}
		}
		return returnImage;
	}
	
	public static BufferedImage RandomDilate(BufferedImage image, Double percentageAmount)
	{
		byte[][] se = new byte[3][3];
		for (int i = 0; i < 3; i++)
		{
			for (int j = 0; j < 3; j++)
				se[i][j] = 1;
		}
		Integer timesToDilate = (int) (randomGen.nextDouble() * Math.min(image.getWidth(), image.getHeight()) * percentageAmount);
		BufferedImage returnImage = image;
		for (int i = 0; i < timesToDilate; i++)
			returnImage = ImageArrayUtil.DilateImage(image, se);
		return returnImage;
	}
	
	public static BufferedImage RandomErode(BufferedImage image, Double percentageAmount)
	{
		byte[][] se = new byte[3][3];
		for (int i = 0; i < 3; i++)
		{
			for (int j = 0; j < 3; j++)
				se[i][j] = 1;
		}
		Integer timesToDilate = (int) (randomGen.nextDouble() * Math.min(image.getWidth(), image.getHeight()) * percentageAmount);
		BufferedImage returnImage = image;
		for (int i = 0; i < timesToDilate; i++)
			returnImage = ImageArrayUtil.ErodeImage(image, se);
		return returnImage;
	}
	
	public static BufferedImage RandomRotate(BufferedImage image, Double percentageRotation)
	{
		Double rotationAngle = randomGen.nextGaussian() * 2 * Math.PI * percentageRotation;
		return ImageArrayUtil.Rotate(image, rotationAngle);
	}
	
}
